import pandas as pd
import os
from glob import glob
import shutil
import dlib
import cv2
import imutils
from imutils import face_utils

df = pd.read_csv('./data/raw.csv')
detector1 = dlib.get_frontal_face_detector()
detector2 = dlib.cnn_face_detection_model_v1('data/mmod_human_face_detector.dat')
count = 0
all_count = df.shape[0]

for i, row in df.iterrows():
    count += 1
    if count % 100 == 0:
        print(count, '/', all_count)
    image_path = row.RawPath
    # 载入图片，检测人脸
    image = cv2.imread(image_path)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    rects = detector1(gray, 2)
    if len(rects) == 0:
        # 如果之前的detector没检测到，就换一个检测
        rects = detector2(gray, 2)
    # 找最大人脸
    max_area = 0
    max_candidate = None
    for (i, rect) in enumerate(rects):
        if hasattr(rect, 'rect'):
            rect = rect.rect
        (x, y, w, h) = face_utils.rect_to_bb(rect)
        if w * h > max_area:
            max_area = w * h
            max_candidate = (x, y, w, h)
    if max_candidate is None:
        print(image_path)
        continue
    x, y, w, h = max_candidate
    x -= 10
    y -= 10
    w += 20
    h += 20
    if x < 0:
        x = 0
    if y < 0:
        y = 0
    cropped = image[y:y + h, x:x + w]
    cropped = imutils.resize(cropped, 128, 128)
    new_image_path = image_path.replace('raw', 'cropped')
    cv2.imwrite(new_image_path, cropped)


new_list = [x.replace('raw', 'cropped') for x in df.RawPath.tolist()]
for i in range(len(new_list)):
    if os.path.exists(new_list[i]) is False:
        new_list[i] = None
df['CroppedPath'] = new_list

df.to_csv('./data/info.csv', index=False)
